«
»

Memoizing polymorphic functions – part two

12th June 2009, 02:04 pm

Part one of this series introduced the problem of memoizing functions involving polymorphic recursion. The caching data structures used in memoization typically handle only one type of argument at a time. For instance, one can have finite maps of differing types, but each concrete finite map holds just one type of key and one type of value.

I extended memoization to handle polymorphic recursion by using an existential type together with a reified type of types. This extension works (afaik), but it is restricted to a particular form for the type of the polymorphic function being memoized, namely

-- Polymorphic function
type k :--> v = forall a. HasType a => k a -> v a

My motivating example is a GADT-based representation of typed lambda calculus, and some of the functions I want to memoize do not fit the pattern. After writing part one, I fooled around and found that I could transform these awkwardly typed polymorphic functions into isomorphic form that does indeed fit the restricted pattern of polymorphic types I can handle.

Awkward types

The first awkwardly typed memoizee is the function application constructor:

type AT = forall a b . (HasType a, HasType b) => E (a -> b) -> E a -> E b

(:^) :: AT

Right away AT misses the required form. It has two HasType constraints, and the first argument is parameterized over two type variables instead of one. However, the second argument looks more promising, so let’s flip the arguments to get an isomorphic type:

forall a b . (HasType a, HasType b) => E a -> E (a -> b) -> E b

And then move the quantifier and constraint on b inside the outer (first) ->:

forall a . HasType a => E a -> (forall b. HasType b => E (a -> b) -> E b)

We’re getting closer. Next, define a newtype wrapper.

newtype A2 a = A2 (forall b. HasType b => E (a -> b) -> E b)

So that AT is isomorphic to

forall a . HasType a => E a -> A2 a

i.e.,

E :--> A2

The function inside of A2 doesn’t have the required form, but another newtype wrapper finishes the job.

newtype EF a b = EF {unEF :: E (a -> b) }

type H' a = EF a :--> E

newtype H a = H { unH :: H' a }

The AT type is isomorphic AP where

type AP = E :--> H

Curried memoization

A “curried memo function” is one that takes one argument and produces another memo function. For a simple memo function, not involving polymorphic recursion, there’s a simple recipe for curried memoization:

memo2 :: (a -> b -> c) -> (a -> b -> c)
memo2 f = memo (memo . f)

Our more polymorphic memo makes currying a little more awkward. First, here’s a helper function for working inside of the representation of an H:

inH :: (H' a -> H' a) -> (H a -> H a)
inH h z = H (h (unH z))

The following more elegant definition doesn’t type-check, due to the rank 2 polymorphism:

inH f = H . f . unH  -- type error

Now our AP memoizer is much like memo2:

memoAP :: AP -> AP
memoAP app' = memo (inH memo . app')

(A more general, consistent type for memoAP is (f :--> H) -> (f :--> H).)

Isomorphisms

Now, to define the isomorphisms. Define

toAP   :: AT -> AP
fromAP :: AP -> AT

The definitions:

toAP app ea = H $  (EF eab) -> app eab ea
fromAP app' eab ea = unH (app' ea) (EF eab)

If you erase the newtype wrappers & unwrappers, you’ll see that toAP and fromAP are both just flip.

I constructed fromAP from the following specification:

toAP (fromAP app') == app'

Transforming step-by-step into equivalent specifications:

 ea -> H $  (EF eab) -> (fromAP app') eab ea == app'

H $  (EF eab) -> (fromAP app') eab ea == app' ea

 (EF eab) -> (fromAP app') eab ea == unH (app' ea)

(fromAP app') eab ea == unH (app' ea) (EF eab)

fromAP app' eab ea == unH (app' ea) (EF eab)

Memoizing vis isomorphisms

Finally, I can memoize

memoAT :: AT -> AT
memoAT app = fromAP (memoAP (toAP app))

Again, a more elegant definition via (.) fails to type-check, due to rank 2 polymorphism.

The Lam (lambda abstraction) constructor can be handled similarly:

Lam  :: (HasType a, HasType b) => V a -> E b -> E (a -> b)

This time, no flip is needed.

I wonder

How far does this isomorphism trick go?

Is there an easier way to memoize polymorphic functions?

Tags: ,  |  Comment (RSS)  |  Trackback

2 Comments

  1. Neel Krishnaswami:

    Hi Conal, I don’t know if this is enough for your purpose, but there’s a beautiful trick I learned from Janis Voigtlander’s 2009 POPL paper “Bidirectionalization for Free”. The idea is to to exploit parametricity — if you have a function that operates over lists parametrically, you can find out what it does on any argument by passing it a list of unique integers, and looking at how it permutes those integers.

    The same trick ought to work for memoization. If we want to write a function

         memo' : (all a. k a -> v a) -> (all a. k a -> v a)`

    then it’s sufficient if we can write the memo function for the Int case. More abstactly, we can do this if k and v are functors, and k also has a way to distribute over the state monad (that is, we want a function seq : (k (state a)) -> (state (k a)). Then, we want to write something like (I’ll use a mutant hybrid ML/Haskell syntax here):

    type state r a = (Int, Int -> r) -> ((Int, Int -> r), a)

eat : r -> state r int 
eat r (size, mapping) = ((size+1, n. if n = size then r else mapping n), size)

run : state r a -> int -> r -> ((int, int -> r), a)
run state initsize initmap = state (initsize, initmap)

seq : k (state r a) -> state r (k a)
seq blah = whatever

memo' (f : all a. k a -> v a) = 
  let table : ref (k int -> option (v int)) =  ref (x. Nothing)   -- allocate a table
  let memofun x : k b = 
    let leaves : k (state b int) = map_k eat x 
    let traversal : state b (k int) = seq leaves 
    let ((size, indices), shape) = run traversal 0 (x. bot)
    case !table shape of 
      Just result_shape -> fmap indices result_shape
    | Nothing -> let result_shape = f shape in 
                 begin
                    table := update table shape result_shape;  -- imperatively update the memo table
                    return (fmap indices result_shape)
                 end

    The idea is that we use map and seq to construct an action that will walk over the key, replacing each polymorphic value with a unique integer index. Map crawls over the term, and we stick a little command in a state monad to eat the value there. Then, we use seq to sequence all those terms, ensuring we a) get unique indices at all the leaves, and b) we can get a mapping from those indices back to the original values.

    Then, we can use memoize the int version of the value, and its result. Given an integer-valued result, we can use the mapping from indices to values to put back the correct values. So v needs to be a functor, and k needs to be a functor with a sequencing distribution property — the traversable stuff ought to be the right general property to ask for.

    18 June 2009, 11:30 am

  2. Conal Elliott » Blog Archive » Memoizing polymorphic functions via unmemoization:

    […] combination of patricia trees (for integer-keyed maps), stable names, and type equality proofs. The second post showed how to transform some functions that do not quite fit the required form so that they do […]

    2 October 2010, 7:41 am

Leave a comment